mod opcode;
mod payload;

pub use opcode::WebSocketOpCode;
pub use payload::WebSocketPayload;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};

use super::{WebSocketParseError, WebSocketReceiveError, WebSocketSendError};

/*

   0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 
  +-+-+-+-+-------+-+-------------+-------------------------------+
  |F|R|R|R| opcode|M| Payload len |    Extended payload length    |
  |I|S|S|S|  (4)  |A|     (7)     |          (16/64)              |
  |N|V|V|V|       |S|             |   (if Payload len==126/127)   |
  | |1|2|3|       |K|             |                               |
  +-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +
  |   Masking-key (if MASK set)   |     (32 bits)                 |
  +-------------------------------+-------------------------------+
  |          Payload Data         | (Length as above)             |
  +-------------------------------+-------------------------------+

*/

#[derive(Debug, Clone)]
pub struct WebSocketFrame {
    pub fin: bool,
    pub opcode: WebSocketOpCode,
    pub payload: WebSocketPayload,
    pub masked: bool,
    pub mask_key: Option<[u8; 4]>,
}

impl WebSocketFrame {
    fn new(opcode: WebSocketOpCode, payload: WebSocketPayload, masked: bool) -> Self {
        let mask_key = if masked {
            Some(rand::random::<[u8; 4]>())
        } else {
            None
        };
        Self {
            fin: true,
            opcode,
            payload,
            masked,
            mask_key,
        }
    }

    pub fn server_text(text: impl Into<String>) -> Self {
        Self::new(WebSocketOpCode::Text, WebSocketPayload::Text(text.into()), false)
    }

    pub fn client_text(text: impl Into<String>) -> Self {
        Self::new(WebSocketOpCode::Text, WebSocketPayload::Text(text.into()), true)
    }

    pub fn server_binary(bytes: impl Into<Vec<u8>>) -> Self {
        Self::new(WebSocketOpCode::Binary, WebSocketPayload::Binary(bytes.into()), false)
    }

    pub fn client_binary(bytes: impl Into<Vec<u8>>) -> Self {
        Self::new(WebSocketOpCode::Binary, WebSocketPayload::Binary(bytes.into()), true)
    }

    pub fn server_close() -> Self {
        Self::new(WebSocketOpCode::Close, WebSocketPayload::Close, false)
    }

    pub fn client_close() -> Self {
        Self::new(WebSocketOpCode::Close, WebSocketPayload::Close, true)
    }

    pub fn server_ping() -> Self {
        Self::new(WebSocketOpCode::Ping, WebSocketPayload::Ping, false)
    }

    pub fn client_ping() -> Self {
        Self::new(WebSocketOpCode::Ping, WebSocketPayload::Ping, true)
    }

    pub fn server_pong() -> Self {
        Self::new(WebSocketOpCode::Pong, WebSocketPayload::Pong, false)
    }

    pub fn client_pong() -> Self {
        Self::new(WebSocketOpCode::Pong, WebSocketPayload::Pong, true)
    }

    pub fn as_str(&self) -> Result<&str, WebSocketParseError> {
        match &self.payload {
            WebSocketPayload::Text(s) => Ok(s.as_str()),
            _ => Err(WebSocketParseError::OpCodeUnmatch("Not a text frame".into())),
        }
    }

    pub fn as_bytes(&self) -> Result<&[u8], WebSocketParseError> {
        match &self.payload {
            WebSocketPayload::Binary(b) => Ok(&b),
            _ => Err(WebSocketParseError::OpCodeUnmatch("Not a text frame".into())),
        }
    }

    pub fn into_string(&self) -> Result<String, WebSocketParseError> {
        Ok(self.as_str()?.to_string())
    }

    pub fn into_bytes(&self) -> Result<Vec<u8>, WebSocketParseError> {
        Ok(self.as_bytes()?.to_vec())
    }

    pub fn encode(&self) -> Vec<u8> {
        let mut buf = Vec::new();
        let mut payload = self.payload.to_bytes();

        if let Some(key) = self.mask_key {
            for i in 0..payload.len() {
                payload[i] ^= key[i % 4];
            }
        }

        let b0 = (if self.fin { 0x80 } else { 0x00 }) | self.opcode.to_u8();
        buf.push(b0);

        let masked = if self.mask_key.is_some() { 0x80 } else { 0 };
        let len = payload.len();
        if len <= 125 {
            buf.push(len as u8 | masked);
        } else if len <= 65535 {
            buf.push(126 | masked);
            buf.extend_from_slice(&(len as u16).to_be_bytes());
        } else {
            buf.push(127 | masked);
            buf.extend_from_slice(&(len as u64).to_be_bytes());
        }

        if let Some(key) = self.mask_key {
            buf.extend_from_slice(&key);
        }

        buf.extend_from_slice(&payload);
        buf
    }

    pub async fn read_from<S>(stream: &mut S) -> Result<Option<Self>, WebSocketReceiveError>
    where
        S: AsyncRead + Unpin,
    {
        let mut header = [0u8; 2];
        match stream.read_exact(&mut header).await {
            Ok(_) => {}
            Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None),
            Err(e) => return Err(e.into()),
        }

        let fin = header[0] & 0x80 != 0;
        let opcode = WebSocketOpCode::from_u8(header[0] & 0x0F);
        let masked = header[1] & 0x80 != 0;
        let mut payload_len = (header[1] & 0x7F) as u64;

        if payload_len == 126 {
            let mut ext = [0u8; 2];
            stream.read_exact(&mut ext).await?;
            payload_len = u16::from_be_bytes(ext) as u64;
        } else if payload_len == 127 {
            let mut ext = [0u8; 8];
            stream.read_exact(&mut ext).await?;
            payload_len = u64::from_be_bytes(ext);
        }

        let mask_key = if masked {
            let mut key = [0u8; 4];
            stream.read_exact(&mut key).await?;
            Some(key)
        } else {
            None
        };

        let mut payload = vec![0u8; payload_len as usize];
        stream.read_exact(&mut payload).await?;

        if let Some(key) = mask_key {
            for i in 0..payload.len() {
                payload[i] ^= key[i % 4];
            }
        }

        let parsed_payload = WebSocketPayload::from(opcode, payload)?;

        Ok(Some(Self {
            fin,
            opcode,
            payload: parsed_payload,
            masked,
            mask_key,
        }))
    }

    pub async fn send_to<S>(&self, stream: &mut S) -> Result<(), WebSocketSendError>
    where
        S: AsyncWrite + Unpin,
    {
        let buf = self.encode();
        stream.write_all(&buf).await?;
        Ok(())
    }
}